/*
 * Copyright (c) 2023 Shenzhen Kaihong Digital Industry Development Co., Ltd.
 * Licensed under the Apache License, Version 2.0 (the "License");
 * you may not use this file except in compliance with the License.
 * You may obtain a copy of the License at
 *
 *     http://www.apache.org/licenses/LICENSE-2.0
 *
 * Unless required by applicable law or agreed to in writing, software
 * distributed under the License is distributed on an "AS IS" BASIS,
 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 * See the License for the specific language governing permissions and
 * limitations under the License.
 */

#include "rtmp_protocol.h"
#include <cstring>
#include <netinet/in.h>
#include <stdexcept>
#include "common/sharing_log.h"
#include "utils.h"

#define C1_DIGEST_SIZE    32
#define C1_KEY_SIZE       128
#define C1_SCHEMA_SIZE    764
#define C1_HANDSHARK_SIZE (RANDOM_LEN + 8)
#define C1_FPKEY_SIZE     30
#define S1_FMS_KEY_SIZE   36
#define S2_FMS_KEY_SIZE   68
#define C1_OFFSET_SIZE    4

namespace OHOS {
namespace Sharing {

RtmpProtocol::RtmpProtocol() {}

RtmpProtocol::~RtmpProtocol()
{
    Reset();
}

void RtmpProtocol::Reset()
{
    chunkSizeIn_ = DEFAULT_CHUNK_LEN;
    chunkSizeOut_ = DEFAULT_CHUNK_LEN;
    bytesSent_ = 0;
    bytesSentLast_ = 0;
    windowsSize_ = 0;
    bandwidth_ = 2500000;
    bandLimitType_ = 2;
    mapChunkData_.clear();
    nowStreamIndex_ = 0;
    nowChunkId_ = 0;
    sendReqId_ = 0;

    streamIndex_ = STREAM_CONTROL;
    nextStepFunc_ = nullptr;
}

void RtmpProtocol::SendAcknowledgement(uint32_t size)
{
    size = htonl(size);
    std::string acknowledgement((char *)&size, 4);
    SendRequest(RTMPMessageType::ACKNOWLEDGEMENT, acknowledgement);
}

void RtmpProtocol::SendAcknowledgementSize(uint32_t size)
{
    size = htonl(size);
    std::string setWindowSize((char *)&size, 4);
    SendRequest(RTMPMessageType::WINDOW_ACK_SIZE, setWindowSize);
}

void RtmpProtocol::SendPeerBandwidth(uint32_t size)
{
    size = htonl(size);
    std::string setPeerBandwidth((char *)&size, 4);
    setPeerBandwidth.push_back((char)0x02);
    SendRequest(RTMPMessageType::SET_PEER_BW, setPeerBandwidth);
}

void RtmpProtocol::SendChunkSize(uint32_t size)
{
    uint32_t len = htonl(size);
    std::string setChunk((char *)&len, 4);
    SendRequest(RTMPMessageType::SET_CHUNK_SIZE, setChunk);
    chunkSizeOut_ = size;
}

void RtmpProtocol::SendPingRequest(uint32_t stamp)
{
    SendUserControl(CONTROL_PING_REQUEST, stamp);
}

void RtmpProtocol::SendPingResponse(uint32_t timestamp)
{
    SendUserControl(CONTROL_PING_RESPONSE, timestamp);
}

void RtmpProtocol::SendSetBufferLength(uint32_t streamId, uint32_t len)
{
    std::string control;
    streamId = htonl(streamId);
    control.append((char *)&streamId, 4);

    len = htonl(len);
    control.append((char *)&len, 4);
    SendUserControl(CONTROL_SETBUFFER, control);
}

void RtmpProtocol::SendUserControl(uint16_t eventType, uint32_t eventData)
{
    std::string control;
    eventType = htons(eventType);
    control.append((char *)&eventType, 2);

    eventData = htonl(eventData);
    control.append((char *)&eventData, 4);
    SendRequest(RTMPMessageType::USER_CONTROL, control);
}

void RtmpProtocol::SendUserControl(uint16_t eventType, const std::string &eventData)
{
    std::string control;
    eventType = htons(eventType);
    control.append((char *)&eventType, 2);
    control.append(eventData);
    SendRequest(RTMPMessageType::USER_CONTROL, control);
}

void RtmpProtocol::SendResponse(RTMPMessageType type, const std::string &str)
{
    if (!dataStarted_ && (type == RTMPMessageType::DATA)) {
        dataStarted_ = true;
    }
    SendRtmp(type, nowStreamIndex_, str, 0, dataStarted_ ? CHUNK_CLIENT_REQUEST_AFTER : CHUNK_CLIENT_REQUEST_BEFORE,
             [this](const char *p, size_t len) { Send(p, len); });
}

void RtmpProtocol::SendInvoke(const std::string &cmd, const AmfValue &val)
{
    AMFEncoder enc;
    enc << cmd << ++sendReqId_ << val;
    SendRequest(RTMPMessageType::CMD, enc.Data());
}

void RtmpProtocol::SendRequest(RTMPMessageType cmd, const std::string &str)
{
    if (cmd <= RTMPMessageType::SET_PEER_BW) {
        SendRtmp(cmd, streamIndex_, str, 0, CHUNK_NETWORK, [this](const char *p, size_t len) { Send(p, len); });
    } else {
        SendRtmp(cmd, streamIndex_, str, 0, CHUNK_SYSTEM, [this](const char *p, size_t len) { Send(p, len); });
    }
}

void RtmpProtocol::SendRtmp(RTMPMessageType type, uint32_t streamId, const std::string &buffer, uint32_t stamp,
                            int chunkId, DataCallBack callback)
{
    DataBuffer::Ptr data = std::make_shared<DataBuffer>();
    data->Assign(buffer.data(), buffer.size());
    SendRtmp(type, streamId, data, stamp, chunkId, callback);
}

void RtmpProtocol::SendRtmp(RTMPMessageType type, uint32_t streamId, const DataBuffer::Ptr &buf, uint32_t stamp,
                            int chunkId, DataCallBack callback)
{
    if (chunkId < 2 || chunkId > 63) {
        SHARING_LOGE("unsupported chunk id");
        return;
    }
    bool ext_stamp = stamp >= 0xFFFFFF;

    DataBuffer::Ptr packet = std::make_shared<DataBuffer>();

    RtmpHeader header{};
    header.fmt = 0;
    header.chunkId = chunkId;
    header.typeId = (uint8_t)type;
    SetBE24(header.timestamp, ext_stamp ? 0xFFFFFF : stamp);
    SetBE24(header.messageLength, (uint32_t)buf->Size());
    SetLE32(header.streamId, streamId);

    packet->Assign((char *)&header, sizeof(RtmpHeader));

    char ext_timestamp[4] = {};
    if (ext_stamp) {
        SetBE32(ext_timestamp, stamp);
    }

    char flag;
    auto flagHeader = (RtmpHeader *)&flag;
    flagHeader->fmt = 3;
    flagHeader->chunkId = chunkId;

    size_t offset = 0;
    size_t totalSize = sizeof(RtmpHeader);
    while (offset < (size_t)buf->Size()) {
        if (offset) {
            packet->Append(&flag, 1);
            totalSize += 1;
        }

        if (ext_stamp) {
            packet->Append(ext_timestamp, 4);
            totalSize += 4;
        }

        size_t chunk = std::min((size_t)DATA_CHUNK_SIZE_OUT, buf->Size() - offset);
        packet->Append(buf->Peek() + offset, (int)chunk);

        totalSize += chunk;
        offset += chunk;
    }

    callback(packet->Peek(), packet->Size());
    SHARING_LOGD("Send RTMP packet, size(%{public}d)\n", packet->Size());
}

// Send C0+C1 to server, Recv S0+S1+S2 from server
void RtmpProtocol::StartClientSession(const std::function<void()> &func)
{
    SHARING_LOGD("trace.");
    DataBuffer::Ptr c0c1 = std::make_shared<DataBuffer>();
    char c0 = HANDSHAKE_PLAINTEXT; // 1 Byte 0x03 for C0
    c0c1->Assign(c0);
    RtmpHandshake c1(0); // 1536 bytes for C1
    c0c1->Append((char *)&c1, sizeof(RtmpHandshake));

    SHARING_LOGD("Send C0+C1 to server");

    nextStepFunc_ = [this, func](const char *data, size_t len) {
        // wait S0+S1+S2
        return HandleS0S1S2(data, len, func);
    };
    Send((char *)c0c1->Peek(), c0c1->Size());
}

// Recv S0+S1+S2, Just check S0 == C0, Send S1 as C2
const char *RtmpProtocol::HandleS0S1S2(const char *data, size_t len, const std::function<void()> &func)
{
    SHARING_LOGD("trace.");
    static DataBuffer::Ptr cache = std::make_shared<DataBuffer>();
    if (len < 1 + 2 * C1_HANDSHARK_SIZE) {
        SHARING_LOGI("data too short\n");
        if (cache->Size() == 0 && data[0] != HANDSHAKE_PLAINTEXT) {
            SHARING_LOGI("only plaintext[0x03] handshake supported");
            return nullptr;
        }

        cache->Append(data, len);

        if (cache->Size() < 1 + 2 * C1_HANDSHARK_SIZE) {
            SHARING_LOGI("data too short return %{public}d", cache->Size());
            return nullptr;
        }

        SHARING_LOGI("data is enough\n");
    } else {
        SHARING_LOGI("data is enough\n");
        cache->Assign(data, len);
    }

    if (cache->Data()[0] != HANDSHAKE_PLAINTEXT) {
        SHARING_LOGE("Invalid S0S1S2 data.");
        return nullptr;
    }

    SHARING_LOGI("send C2 (same as s1)\n");
    Send((char *)data + 1, sizeof(RtmpHandshake));
    cache->Clear();
    nextStepFunc_ = [this](const char *data, size_t len) { return HandleRtmp(data, len); };

    func();

    return data + 1 + 2 * C1_HANDSHARK_SIZE;
}

static constexpr size_t HEADER_LENGTH[] = {12, 8, 4, 1};

const char *RtmpProtocol::HandleRtmp(const char *data, size_t size)
{
    SHARING_LOGD("trace");
    if (!cache_) {
        cache_ = std::make_shared<DataBuffer>();
    }

    SHARING_LOGD("remain size: %{public}d, incoming size: %{public}d", cache_->Size(), size);
    cache_->Append(data, size);

    auto ptr = cache_->Peek();
    auto end = cache_->Peek() + cache_->Size();
    size_t len = cache_->Size();
    while (len) {
        size_t offset = 0;
        auto header = (RtmpHeader *)ptr;
        auto header_len = HEADER_LENGTH[header->fmt];

        nowChunkId_ = header->chunkId;
        switch (nowChunkId_) {
            case 0: {
                if (len < 2) {
                    SHARING_LOGI("need more data");
                    break;
                }
                nowChunkId_ = 64 + (uint8_t)(ptr[1]);
                offset = 1;
                break;
            }

            case 1: {
                if (len < 3) {
                    SHARING_LOGI("need more data");
                    break;
                }
                nowChunkId_ = 64 + ((uint8_t)(ptr[2]) << 8) + (uint8_t)(ptr[1]);
                offset = 2;
                break;
            }

            default:
                break;
        }

        if (len < header_len + offset) {
            SHARING_LOGI("need more data");
            break;
        }

        header = (RtmpHeader *)(ptr + offset);
        auto &pr = mapChunkData_[nowChunkId_];

        auto &now_packet = pr.first;
        auto &last_packet = pr.second;

        if (!now_packet) {
            now_packet = RtmpPacket::Create();
            if (last_packet) {
                *now_packet = *last_packet;
                now_packet->Clear();
            }
            now_packet->isAbsStamp = false;
        }

        auto &chunk_data = *now_packet;
        chunk_data.chunkId = nowChunkId_;

        switch (header_len) {
            case 12:
                chunk_data.isAbsStamp = true;
                chunk_data.streamId = *(uint32_t *)(header->streamId);
                chunk_data.bodySize = LoadBE24(header->messageLength);
                chunk_data.typeId = (RTMPMessageType)header->typeId;
                break;
            case 8:
                chunk_data.bodySize = LoadBE24(header->messageLength);
                chunk_data.typeId = (RTMPMessageType)header->typeId;
                chunk_data.tsField = LoadBE24(header->timestamp);
                break;
            case 4:
                chunk_data.tsField = LoadBE24(header->timestamp);
                break;
            default:
                break;
        }

        auto timestamp = chunk_data.tsField;
        if (chunk_data.tsField == 0xFFFFFF) {
            if (len < header_len + offset + 4) {
                SHARING_LOGI("need more data");
                break;
            }
            timestamp = LoadBE32((uint8_t *)ptr + offset + header_len);
            offset += 4;
        }

        auto more = std::min(chunkSizeIn_, (size_t)(chunk_data.bodySize /* - chunk_data.Size()*/));
        if (len < header_len + offset + more) {
            SHARING_LOGI("need more data");
            break;
        }
        if (more) {
            chunk_data.Assign(ptr + header_len + offset, (int)more);
        }
        ptr += header_len + offset + more;
        len -= header_len + offset + more;
        if ((size_t)chunk_data.Size() == chunk_data.bodySize) {
            nowStreamIndex_ = (int)chunk_data.streamId;
            chunk_data.timestamp = timestamp + (chunk_data.isAbsStamp ? 0 : chunk_data.timestamp);
            last_packet = now_packet;

            if (chunk_data.bodySize) {
                HandleRtmpChunk(now_packet);
            } else {
                now_packet = nullptr;
            }
        }
    }

    if (ptr < end) {
        DataBuffer::Ptr cacheNew = std::make_shared<DataBuffer>();
        cacheNew->Assign(ptr, end - ptr);
        cache_ = cacheNew;
        SHARING_LOGD("remain: %{public}d, use new cache", (int)(end - ptr));
    } else {
        cache_->Clear();
        SHARING_LOGD("Clear cache");
    }

    return ptr;
}

void RtmpProtocol::HandleRtmpChunk(RtmpPacket::Ptr packet)
{
    SHARING_LOGD("trace");
    auto &chunk_data = *packet;
    switch (chunk_data.typeId) {
        case RTMPMessageType::ACKNOWLEDGEMENT: {
            if (chunk_data.Size() < 4) {
                SHARING_LOGD("MSG_ACK: Not enough data");
                return;
            }
            auto bytePeerRecv = LoadBE32(chunk_data.Data());

            SHARING_LOGD("-> ACKNOWLEDGEMENT: %{public}d, %{public}02x-%{public}02x-%{public}02x-%{public}02x\n",
                         bytePeerRecv, chunk_data.Data()[0], chunk_data.Data()[1], chunk_data.Data()[2],
                         chunk_data.Data()[3]);
            break;
        }

        case RTMPMessageType::SET_CHUNK_SIZE: {
            if (chunk_data.Size() < 4) {
                SHARING_LOGD("MSG_SET_CHUNK :Not enough data");
                return;
            }
            chunkSizeIn_ = LoadBE32(chunk_data.Data());
            SHARING_LOGD("-> SET_CHUNK_SIZE: %{public}zu\n", chunkSizeIn_);
            break;
        }

        case RTMPMessageType::USER_CONTROL: {
            // user control message
            if (chunk_data.Size() < 2) {
                SHARING_LOGD("MSG_USER_CONTROL: Not enough data.");
                return;
            }
            uint16_t eventType = LoadBE16(chunk_data.Data());

            switch (eventType) {
                case CONTROL_PING_REQUEST: {
                    if (chunk_data.Size() - 2 < 4) {
                        SHARING_LOGD("CONTROL_PING_REQUEST: Not enough data.");
                        return;
                    }
                    uint32_t timeStamp = LoadBE32(chunk_data.Data() + 2);
                    SHARING_LOGD("CONTROL_PING_REQUEST: %{public}d", timeStamp);
                    SendUserControl(CONTROL_PING_RESPONSE, timeStamp);
                    break;
                }

                case CONTROL_PING_RESPONSE: {
                    if (chunk_data.Size() - 2 < 4) {
                        SHARING_LOGD("CONTROL_PING_RESPONSE: Not enough data.");
                        return;
                    }
                    uint32_t timeStamp = LoadBE32(chunk_data.Data() + 2);
                    SHARING_LOGD("CONTROL_PING_REQUEST: %{public}d", timeStamp);
                    break;
                }

                case CONTROL_STREAM_BEGIN: {
                    if (chunk_data.Size() - 2 < 4) {
                        SHARING_LOGD("CONTROL_STREAM_BEGIN: Not enough data: %{public}d", chunk_data.Size() - 2);
                        break;
                    }
                    uint32_t streamId = LoadBE32(chunk_data.Data() + 2);
                    OnStreamBegin(streamId);
                    SHARING_LOGD("CONTROL_STREAM_BEGIN: %{public}d", streamId);
                    break;
                }

                case CONTROL_STREAM_EOF: {
                    if (chunk_data.Size() - 2 < 4) {
                        SHARING_LOGD("CONTROL_STREAM_EOF: Not enough data.");
                        return;
                    }
                    uint32_t streamId = LoadBE32(chunk_data.Data() + 2);
                    OnStreamEof(streamId);
                    SHARING_LOGD("CONTROL_STREAM_EOF: %{public}d", streamId);
                    break;
                }

                case CONTROL_STREAM_DRY: {
                    if (chunk_data.Size() - 2 < 4) {
                        SHARING_LOGD("CONTROL_STREAM_DRY: Not enough data.");
                        return;
                    }
                    uint32_t streamId = LoadBE32(chunk_data.Data() + 2);
                    OnStreamDry(streamId);
                    SHARING_LOGD("CONTROL_STREAM_DRY: %{public}d", streamId);
                    break;
                }

                default:
                    SHARING_LOGD("unhandled user control: %{public}d", eventType);
                    break;
            }
            break;
        }

        case RTMPMessageType::WINDOW_ACK_SIZE: {
            windowsSize_ = LoadBE32(chunk_data.Data());
            SHARING_LOGD("-> WINDOW_ACK_SIZE: %{public}d, %{public}02x-%{public}02x-%{public}02x-%{public}02x\n",
                         windowsSize_, chunk_data.Data()[0], chunk_data.Data()[1], chunk_data.Data()[2],
                         chunk_data.Data()[3]);
            break;
        }

        case RTMPMessageType::SET_PEER_BW: {
            bandwidth_ = LoadBE32(chunk_data.Data());
            bandLimitType_ = chunk_data.Data()[4];
            SHARING_LOGD("-> SET_PEER_BW: %{public}d-%{public}d", bandwidth_, (int)bandLimitType_);
            break;
        }

        case RTMPMessageType::AGGRAGATE: {
            auto ptr = chunk_data.Data();
            auto ptr_tail = ptr + chunk_data.Size();
            uint32_t latest_ts, timestamp;
            timestamp = chunk_data.timestamp;
            bool first_message = true;
            while (ptr + 8 + 3 < ptr_tail) {
                auto type = *ptr;
                ptr += 1;
                auto size = LoadBE32(ptr);
                ptr += 3;
                auto ts = LoadBE24(ptr);
                ptr += 3;
                ts |= (*ptr << 24);
                ptr += 1;
                ptr += 3;
                size += 4;
                if (ptr + size > ptr_tail) {
                    break;
                }
                if (!first_message) {
                    timestamp += ts - latest_ts;
                }
                first_message = false;
                latest_ts = ts;
                auto sub_packet_ptr = RtmpPacket::Create();
                auto &sub_packet = *sub_packet_ptr;
                sub_packet.Assign((char *)ptr, size);
                sub_packet.typeId = (RTMPMessageType)type;
                sub_packet.bodySize = size;
                sub_packet.timestamp = timestamp;
                sub_packet.streamId = chunk_data.streamId;
                sub_packet.chunkId = chunk_data.chunkId;
                HandleRtmpChunk(sub_packet_ptr);
                ptr += size;
            }
            break;
        }

        default: {
            bytesRecv_ += packet->Size();
            if (windowsSize_ > 0 && bytesRecv_ - bytesRecvLast_ >= windowsSize_) {
                bytesRecvLast_ = bytesRecv_;
                SendAcknowledgement(bytesRecv_);
            }
            OnRtmpChunk(packet);
            break;
        }
    }
}

} // namespace Sharing
} // namespace OHOS